import concurrent.futures
from collections import defaultdict
from typing import Tuple, Dict

import faiss
import numpy as np

import data_source


class SearchRes(object):
    def __init__(self, id1: int, distance: float):
        self.id = id1
        self.distance = distance

    def to_dict(self):
        return {'id': str(self.id), 'distance': str(self.distance)}

    def __str__(self):
        return f"id: {self.id},distance: {self.distance}"

    def __repr__(self):
        return self.__str__()


def convert_res(batch_size: int, distances: list, ids: list) -> list[list[SearchRes]]:
    res = []
    for i in range(batch_size):
        res.append([SearchRes(id1, distance) for id1, distance in zip(ids[i], distances[i])])
    return res


# 加载文件并转换为numpy数组
def convert_xm_file(path: str):
    print("convert_xm_file")
    ids = []  # batchSize * 1
    vecs = []  # batchSize * dim
    with open(path) as f:
        line = f.readline()
        while len(line) != 0:
            arr = line.split(':')
            ids.append(arr[0])
            vecs.append(arr[1].strip().split(','))
            line = f.readline()
    return np.array(ids).astype('int'), np.array(vecs).astype('float32')


# stage: 0：空, 1：创建, 2：初始化好（train+add）
class FaissIndex(object):

    def __init__(self, name: str, d: int, index_factory: str = "Flat", hdfs_path: str = None,
                 local_vector_path: str = None):
        self.stage = 0
        self.name = name
        self.d = d
        if d <= 0:
            return
        self.index_factory = index_factory
        self.index_f = faiss.index_factory(d, index_factory)
        self.stage = 1
        if local_vector_path is not None and local_vector_path != '':
            print("init load local_vector_path")
            self.local_vector_path = local_vector_path
            self.init_add_local(self.local_vector_path)
            return
        if hdfs_path is None or hdfs_path == '':
            return
        self.local_path = './data/' + name
        self.hdfs_path = hdfs_path
        self.init_add_hdfs()

    def init_add_hdfs(self):
        success = data_source.download(self.hdfs_path, self.local_path)
        if not success:
            print(f"download {self.hdfs_path} failed!")
            return
        ids, vecs = convert_xm_file(self.local_path)
        self.init_add(ids, vecs)

    def init_add_local(self, local_path: str):
        ids, vecs = convert_xm_file(local_path)
        self.init_add(ids, vecs)

    def init_add(self, ids: list, vecs: list):
        print("train index,name=", self.name, "d=", self.d)
        self.index_f.train(vecs)
        self.index_f.add_with_ids(vecs, ids)
        self.stage = 2

    async def add(self, x: list):
        self.index_f.add(x)

    def search(self, x: np.array, k: int) -> list[SearchRes]:
        distances, ids = self.index_f.search(x, k)
        return [SearchRes(id1, distance) for id1, distance in zip(ids[0], distances[0])]

    def multi_search(self, x: list, k: int) -> list[list[SearchRes]]:
        batch_size = len(list) // self.d
        distances, ids = self.index_f.search(x, k)
        return convert_res(batch_size, distances, ids)

    def info(self):
        if self.name == 'init':
            return 'index initializing'
        return {'name': self.name, 'd': self.d,
                'index_trained': self.index_f.is_trained,
                'index_men': 'todo',
                'index_size': self.index_f.ntotal,
                'index_factory': self.index_factory
                }


INDEX_PLACEHOLDER = FaissIndex("init", 0, "Flat", "")
INDEX_EMPTY = FaissIndex("empty", 0, "Flat", "")


def to_np_arr(t: Tuple[list, list]) -> Tuple[np.array, np.array]:
    return np.array(t[0]).astype('int'), np.array(t[1]).astype('float32')


# TODO 支持标签检索
# 文件格式：tag:id:f1,f2,f3,...
class MultiFaissIndex(object):
    # def __init__(self, name: str, d: int, index_factory: str = "Flat", hdfs_path: str = None,
    #              local_vector_path: str = None):

    def init_add_hdfs(self):
        success = data_source.download(self.hdfs_path, self.local_path)
        if not success:
            print(f"download {self.hdfs_path} failed!")
            return
        tag_dict = self.convert_xm_file(self.local_path)
        self.init_add(tag_dict)

    def init_add_local(self, local_path: str):
        ids, vecs = self.convert_xm_file(local_path)
        self.init_add(ids, vecs)

    def init_add(self, tag_dict: Dict[str, Tuple[list, list]]):
        print("train index,name=", self.name, "d=", self.d)

        # self.index_f.train(None)
        # print("add index,name=", self.name, "d=", self.d)
        # self.index_f.add_with_ids(vecs, ids)
        self.stage = 2

    def convert_xm_file(self, path: str) -> Dict[str, Tuple[np.array, np.array]]:
        print("multi convert_xm_file")
        tag_dict: Dict[str, Tuple[list, list]] = defaultdict(lambda: ([], []))
        with open(path) as f:
            line = f.readline()
            while len(line) != 0:
                arr = line.split(':')
                tag = arr[0]
                tag_dict[tag][0].append(arr[1])
                tag_dict[tag][1].append(arr[2].strip().split(','))
                line = f.readline()
        return {key: to_np_arr(value) for key, value in tag_dict.items()}


# 全局索引注册
global_index_register: Dict[str, FaissIndex] = {}
executor = concurrent.futures.ThreadPoolExecutor()

print('模块初始化')


def get_index(name: str) -> FaissIndex:
    return global_index_register.get(name, None)


def set_index(name: str, index: FaissIndex):
    global_index_register[name] = index


def index_register(name: str, dim: int, index_factory: str, hdfs_path: str, local_path: str):
    try:
        idx = FaissIndex(name, dim, index_factory, hdfs_path, local_path)
        global_index_register[name] = idx
        try:
            test_vec = np.array(np.random.random(dim)).reshape(1, -1)
            print('test_vec=', test_vec)
            res = idx.search(test_vec, 2)
            print('test_search_res=', res)
        except Exception as e2:
            print('test_search_err=', e2)
        print('index register success,name=', name, ',global_index_register.size=', len(global_index_register))
    except Exception as e:
        print('index register err,e=', e)


def infos() -> list[dict]:
    print('global_index_register.size=', len(global_index_register))
    return [index.info() for index in global_index_register.values()]


def async_index_register(name: str, dim: int, index_factory: str, hdfs_path: str, local_path: str):
    executor.submit(index_register, name, dim, index_factory, hdfs_path, local_path)


# 初始化加载
def post_construct(name: str, dim: int, index_factory: str, hdfs_path: str = None, local_path: str = None):
    print('post_construct')
    async_index_register(name, dim, index_factory, hdfs_path, local_path)

# if __name__ == '__main__':
#     a = '111222\n'
#     print(len(a))
#     # print(len(a.split('\b')[1]))
#     print(len(a.strip()))
#     # ids = []
#     # for s in a.split("\b"):
#     #     if len(s) > 0:
#     #         ids.append(int(s))
#     #
#     # ids
#     print('-----')
